# -*- coding: utf-8 -*-
"""
Created on Tue Sep  6 11:48:01 2022

@author: wulong
"""
import numpy as np
from casadi import *
from Basic_paras import*
# In[1] Get Day-ahead solution
def get_results_h(optimalValues, pred_horzn):
    
    u = np.zeros((pred_horzn, Nuc_h))
    z = np.zeros((pred_horzn, Nuz_h))
    x = np.zeros((pred_horzn, Nx_h))
    y = np.zeros((pred_horzn, Ny_h))
    e = np.zeros((pred_horzn, 2))
    yeb = np.zeros((pred_horzn, 1))
    
    seed_point = Nuc_h + Nuz_h + Nx_h + Ny_h + 2 + 1
    
    i_u = 0
    i_z = Nuc_h
    i_x = Nuc_h + Nuz_h
    i_y = Nuc_h + Nuz_h + Nx_h
    i_e = Nuc_h + Nuz_h + Nx_h + Ny_h
    i_yeb = Nuc_h + Nuz_h + Nx_h + Ny_h + 2
    
    for j in range(pred_horzn):
        
        u[j,:] = optimalValues[int(i_u + seed_point*j):int(i_u + seed_point*j + Nuc_h)]
        z[j,:] = optimalValues[int(i_z + seed_point*j):int(i_z + seed_point*j + Nuz_h)]
        x[j,:] = optimalValues[int(i_x + seed_point*j):int(i_x + seed_point*j + Nx_h)]
        y[j,:] = optimalValues[int(i_y + seed_point*j):int(i_y + seed_point*j + Ny_h)]
        e[j,:] = optimalValues[int(i_e + seed_point*j):int(i_e + seed_point*j + 2)]
        yeb[j,:] = optimalValues[int(i_yeb + seed_point*j):int(i_yeb + seed_point*j + 1)]
        
        pass
    
    return u, z, x, y, e, yeb

# In[2] Get S-EMPC solution
def get_results_s(optimalValues, pred_horzn):
    
    u = np.zeros((pred_horzn, Nuc_s))
    u_f = np.zeros((pred_horzn, Nuc_f))
    x = np.zeros((pred_horzn, Nx_s))
    x_f = np.zeros((pred_horzn, Nx_f))
    y = np.zeros((pred_horzn, Ny_s))
    yspe = np.zeros((pred_horzn, 1))
    
    seed_point = Nuc_s + Nuc_f + Nx_s + Nx_f + Ny_s + 1
    
    i_u = 0
    i_uf = Nuc_s
    i_x = Nuc_s + Nuc_f
    i_xf = Nuc_s + Nuc_f + Nx_s
    i_y = Nuc_s + Nuc_f + Nx_s + Nx_f
    i_yspe = Nuc_s + Nuc_f + Nx_s + Nx_f + Ny_s
    
    for j in range(pred_horzn):
        
        u[j,:] = optimalValues[int(i_u + seed_point*j):int(i_u + seed_point*j + Nuc_s)]
        u_f[j,:] = optimalValues[int(i_uf + seed_point*j):int(i_uf + seed_point*j + Nuc_f)]
        x[j,:] = optimalValues[int(i_x + seed_point*j):int(i_x + seed_point*j + Nx_s)]
        x_f[j,:] = optimalValues[int(i_xf + seed_point*j):int(i_xf + seed_point*j + Nx_f)]
        y[j,:] = optimalValues[int(i_y + seed_point*j):int(i_y + seed_point*j + Ny_s)]
        yspe[j,:] = optimalValues[int(i_yspe + seed_point*j):int(i_yspe + seed_point*j + 1)]
        
        pass
    
    return u, u_f, x, x_f, y, yspe

# In[3] Get F-EMPC 1 solution
def get_results_f1(optimalValues, pred_horzn):
    
    u = np.zeros((pred_horzn, Nuc_f1))
    x = np.zeros((pred_horzn, Nx_f1))
    y = np.zeros((pred_horzn, Ny_f1))
    e = np.zeros((pred_horzn, 2))
    
    seed_point = Nuc_f1 + Nx_f1 + Ny_f1 + 2
    
    i_u = 0
    i_x = Nuc_f1
    i_y = Nuc_f1 + Nx_f1
    i_e = Nuc_f1 + Nx_f1 + Ny_f1
    
    for j in range(pred_horzn):
        
        u[j,:] = optimalValues[int(i_u + seed_point*j):int(i_u + seed_point*j + Nuc_f1)]
        x[j,:] = optimalValues[int(i_x + seed_point*j):int(i_x + seed_point*j + Nx_f1)]
        y[j,:] = optimalValues[int(i_y + seed_point*j):int(i_y + seed_point*j + Ny_f1)]
        e[j,:] = optimalValues[int(i_e + seed_point*j):int(i_e + seed_point*j + 2)]
        
        pass
    
    return u, x, y, e

# In[4] Get F-EMPC 2 solution
def get_results_f2(optimalValues, pred_horzn):
    
    u = np.zeros((pred_horzn, Nuc_f2))
    x = np.zeros((pred_horzn, Nx_f2))
    y = np.zeros((pred_horzn, Ny_f2))
    e = np.zeros((pred_horzn, 1))
    
    seed_point = Nuc_f2 + Nx_f2 + Ny_f2 + 1
    
    i_u = 0
    i_x = Nuc_f2
    i_y = Nuc_f2 + Nx_f2
    i_e = Nuc_f2 + Nx_f2 + Ny_f2
    
    for j in range(pred_horzn):
        
        u[j,:] = optimalValues[int(i_u + seed_point*j):int(i_u + seed_point*j + Nuc_f2)]
        x[j,:] = optimalValues[int(i_x + seed_point*j):int(i_x + seed_point*j + Nx_f2)]
        y[j,:] = optimalValues[int(i_y + seed_point*j):int(i_y + seed_point*j + Ny_f2)]
        e[j,:] = optimalValues[int(i_e + seed_point*j):int(i_e + seed_point*j + 1)]
        
        pass
    
    return u, x, y, e

# In[5] Get F-EMPC 3 solution
def get_results_f3(optimalValues, pred_horzn):
    
    u = np.zeros((pred_horzn, Nuc_f3))
    x = np.zeros((pred_horzn, Nx_f3))
    y = np.zeros((pred_horzn, Ny_f3))
    e = np.zeros((pred_horzn, 1))
    
    seed_point = Nuc_f3 + Nx_f3 + Ny_f3 + 1
    
    i_u = 0
    i_x = Nuc_f3
    i_y = Nuc_f3 + Nx_f3
    i_e = Nuc_f3 + Nx_f3 + Ny_f3
    
    for j in range(pred_horzn):
        
        u[j,:] = optimalValues[int(i_u + seed_point*j):int(i_u + seed_point*j + Nuc_f3)]
        x[j,:] = optimalValues[int(i_x + seed_point*j):int(i_x + seed_point*j + Nx_f3)]
        y[j,:] = optimalValues[int(i_y + seed_point*j):int(i_y + seed_point*j + Ny_f3)]
        e[j,:] = optimalValues[int(i_e + seed_point*j):int(i_e + seed_point*j + 1)]
        
        pass
    
    return u, x, y, e
